package com.wdemo.service.impl;

import cn.hutool.core.util.ObjectUtil;
import com.wdemo.service.EmbeddingService;
import com.wdemo.utils.FileUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.core.io.Resource;
import org.springframework.scheduling.annotation.Async;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @ClassName EmbeddingServiceImpl
 * @Description 向量数据库实现类
 * @Author WDEMO
 * @Date 2025/8/29 17:22
 */
@Slf4j
public class EmbeddingServiceImpl implements EmbeddingService {


    private VectorStore vectorStore;


    public EmbeddingServiceImpl(VectorStore vectorStore){
        this.vectorStore=vectorStore;
    }

    @Override
    public void saveVectorStore(List<String> messages) {
        log.info("保存向量数据使用的VectorStore: {}", vectorStore.getClass().getSimpleName());
        log.info("保存到向量数据库中，消息数据：{}", messages);
        List<Document> documents = messages.stream()
                .map(message -> new Document(message))
                .collect(Collectors.toList());
        vectorStore.add(documents);
        log.info("保存到向量数据库成功，数量：{}", documents.size());
    }


    @Async
    @Override
    public void saveVectorStore(Resource resource) {
        log.info("保存到向量数据库中，消息数据：{}", resource);
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build()
        );
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = pagePdfDocumentReader.read();

        // 添加元数据
        documents.forEach(document -> {
            document.getMetadata().put("file_name", resource.getFilename());
            document.getMetadata().put("upload_time", System.currentTimeMillis());
        });
        vectorStore.add(documents);

        log.info("保存到向量数据库成功，数量：{}", documents.size());
    }

    @Async
    @Override
    public void saveVectorStore(Resource resource, String filePath) {
        log.info("保存到向量数据库中，消息数据：{}", resource);
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build()
        );
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = pagePdfDocumentReader.read();

        // 添加元数据
        documents.forEach(document -> {
            if (ObjectUtil.isNull(filePath)) {
                document.getMetadata().put("file_name", resource.getFilename());
            } else {
                String fileName = FileUtils.extractFilenameFromOssPath(filePath);
                document.getMetadata().put("file_name", fileName);
            }
            document.getMetadata().put("upload_time", System.currentTimeMillis());
        });
        vectorStore.add(documents);

        log.info("保存到向量数据库成功，数量：{}", documents.size());
    }

    @Override
    public List<Document> searchAll() {
        log.info("查询向量数据使用的VectorStore: {}", vectorStore.getClass().getSimpleName());

        List<Document> documents = vectorStore.similaritySearch(
                SearchRequest.builder()
                        .query("")
                        .topK(999)
                        .build());
        return documents;
    }

    @Override
    public List<Document> search(String message) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        List<Document> documents = vectorStore
                .similaritySearch(SearchRequest.builder()
                        .query(message)
                        .topK(5)
//                        .filterExpression(b.eq("page_number",2).build())
                        .build());

        return documents;
    }
}
